#!/usr/bin/env python
# -*- coding=utf-8 -*-
import os
import argparse
import configparser
import gzip
from collections import defaultdict
from datetime import datetime
from textwrap import dedent
from utils import TitleParser,sample_info_parse,gopen

#读取配置文件
configs = configparser.ConfigParser()
root_dir = os.path.dirname(os.path.abspath(__file__))
configs.read(os.path.join(root_dir,'fastwork.ini'))
print(configs.sections())
ini_tag = 'mark_confidence_tj' if 'nj' not in os.environ['HOSTNAME'] else 'mark_confidence_nj'
#全局变量
pwd = os.path.abspath('.')
today_str = datetime.now().strftime('%Y-%m-%d')

class AnnovarParser():

    def __init__(self,args):
        self.infile = args.get('infile')
        self.indata = gopen(self.infile,'r')
        self.title = self.indata.readline()
        self.title_list = self.title.strip().split('\t')
        self.tp = TitleParser(self.title)
        self.opre = args.get('opre')
        self.add_sample = args.get('mutated_samples')
        self.no_counter = args.get('no_counter')
        self.mapdir = args.get('mapdir')    
        self.confidence = args.get('confidence')
        self.bam2html = args.get('bam2html')
        self.acmg = args.get("acmg")
        self.sample_info = args.get('confidence')
        if self.sample_info:
            self.sample_info_path = os.path.abspath(self.sample_info)

    def tags(self):    
        tags = [] # tagname,titlename,addtagfunction,args
        if self.add_sample and self.no_counter:
            tags.append(('samples','Patients_With_The_Variant',
                self.add_mutated_samples,{"no_counter":self.no_counter}))
        elif self.add_sample:
            tags.append(('samples','Mutated_number\tPatients_With_The_Variant',
                self.add_mutated_samples,{"no_counter":self.no_counter}))
        #添加其他标签的方法是在此处定义函数及标签并将它们添加到tags中,同时添加argparse参数
        #...
        variants_info = self.add_tags(tags)
        return variants_info

    def get_acmg_tags(self):
        """
        get the ACMG classification of variants from ACMG module
        """
        self.sample_info

    def main(self):
        if self.add_sample:
            variants_info = self.tags()
        if self.bam2html or self.confidence:
            if (not self.mapdir) or (not os.path.exists(self.mapdir)):
                raise Exception('--mapdir/-m is required if you want to mark confidence or to use bam2html function')
        
        if self.bam2html:
            if not locals().get("variants_info"):
                variants_info = self.variants_info_only()
            self.generate_command(variants_info)

        if self.confidence:
            bam_list = self.find_bam_list()
            self.make_confidence(bam_list)
        self.indata.close()
        
    def find_bam_list(self):
        start_idx = self.title_list.index('FORMAT') + 1
        end_idx = self.title_list.index('Ori_REF')
        all_samples = [self.title_list[i] for i in range(start_idx,end_idx)]
        bam_list = []
        for sample in all_samples:
            bam_path1 = os.path.join(self.mapdir,sample,sample+'.final.bam')
            bam_path2 = os.path.join(self.mapdir,sample+'.'+sample,sample+'.final.bam')
            if os.path.exists(bam_path1):
                bam_list.append(bam_path1)
            elif os.path.exists(bam_path2):
                bam_list.append(bam_path2)
            else:
                print('Can NOT find final bam by neighter the {} nor {}'.format(
                    bam_path1,bam_path2
                ))
                print('[Warning]Jump this sample!')
                pass
        return bam_list

    def variants_info_only(self):
        """
        """
        variants_info = defaultdict(list)
        for line in self.indata:
            line_list = line.strip().split('\t')
            self.find_mutated_samples(line_list,variants_info)
        
        return variants_info

    def add_tags(self,tags):
        """
        add one or more tags to annovar file
        tags = [(tagname,titlename,addtagfunction,args)]
        """
        variants_info = defaultdict(list)
        #self.tag_ofile 可用于之后的mark confidence
        self.tag_ofile = os.path.join(pwd,self.opre + '.xls')
        with open(self.tag_ofile,'w') as odata:
            title_names = [tag[1] for tag in tags]
            tagnames = [tag[0] for tag in tags]
            odata.write(self.title.rstrip() + "\t" + "\t".join(title_names) + '\n')
            for line in self.indata:
                line_list = line.strip().split('\t')
                if "samples" in tagnames:
                    for tagname,_,func,args in tags:
                        if tagname == "samples":
                            line_list = func(line_list,variants_info,**args)
                        #add tag and func here
                        #...

                odata.write('\t'.join([str(i) for i in line_list])+'\n')
        print('Output final annotation probands filtered file to {ofile}'.format(
            ofile=self.tag_ofile))    
        return variants_info

    def find_mutated_samples(self,line_list,variants_info):
        """
        """
        chrom = self.tp.get_field(line_list,'CHROM')
        pos = self.tp.get_field(line_list,'POS')
        start_idx = self.tp.get_idx('FORMAT') + 1
        end_idx = self.tp.get_idx('Ori_REF')
        mutated_samples = []
        for idx in range(start_idx,end_idx):
            if line_list[idx].split(':')[0] not in ['./.','0/0','.','.|.']:
                sample = self.title_list[idx].strip()
                mutated_samples.append(sample)
                variants_info[sample].append([chrom,pos])
        return mutated_samples

    def add_mutated_samples(self,line_list,variants_info,no_counter=False):
        """
        """
        mutated_samples = self.find_mutated_samples(line_list,variants_info)
        if no_counter:
            line_list.append(",".join(mutated_samples))
        else:
            line_list.append(len(mutated_samples))
            line_list.append(",".join(mutated_samples))
        return line_list

    def generate_command(self,variants_info):
        """
        generate bam2html command
        """
        bam2html_ofile = self.opre + '.sh'
        commands = []
        for sample in variants_info:
            bam_path1 = os.path.join(self.mapdir,sample,sample+'.final.bam')
            bam_path2 = os.path.join(self.mapdir,sample+'.'+sample,sample+'.final.bam')
            if os.path.exists(bam_path1):
                for chrom,pos in variants_info[sample]:
                    command = 'bam2html {bam} {chrom}:{pos}'.format(
                        bam=bam_path1,chrom=chrom,pos=pos,
                    )
                    commands.append(command)

            elif os.path.exists(bam_path2):
                for chrom,pos in variants_info[sample]:
                    command = 'bam2html {bam} {chrom}:{pos}'.format(
                        bam=bam_path2,chrom=chrom,pos=pos
                    )
                    commands.append(command)

            else:
                print('Can NOT find final bam by neighter the {} nor {}'.format(
                    bam_path1,bam_path2
                ))
                print('[Warning]Jump this sample!')
                pass

        with open(bam2html_ofile,'w') as odata:
            for cmd in commands:
                odata.write(cmd + '\n')
        print('Output bam2html commands to {file}'.format(file=bam2html_ofile))


    def make_confidence(self,bam_list):
        """
        """
        mark_conf_dir = os.path.join(pwd,"mark_confidence."+ today_str)
        if not os.path.exists(mark_conf_dir):
            os.mkdir(mark_conf_dir)
        pileup_script =  configs.get(ini_tag,'pileup_script')
        mark_script =  configs.get(ini_tag,'mark_script')
        pileup_file = self.opre + '.pileup'
        middle_shell_name = self.opre + '.make_pileup.sh'
        output_prefix = self.opre
        confidence_shell = os.path.join(mark_conf_dir,'work' + '.confidence_mark.'+ today_str +'.sh')
        bam_str = ",".join(bam_list)
        if self.add_sample:
            candidate_sites_file = os.path.abspath(self.tag_ofile)
        else:
            candidate_sites_file = os.path.abspath(self.infile)

        sample_info = self.sample_info_path
        template = dedent("""\
            #Step1 Generating Pileup File for candidate sites
            python {pileup_script} \\
                -B {bam_str} \\
                -input {candidate_sites_file} \\
                -output {mark_conf_dir}/{pileup_file} \\
                -O {mark_conf_dir}/{middle_shell_name}

            #Step2 Marking Confidence
            python {mark_script} \\
                -I {candidate_sites_file} \\
                -O {mark_conf_dir}/{output_prefix} \\
                -Pileup {mark_conf_dir}/{pileup_file} \\
                -R {sample_info} \\
                -Check Y
            """.format(**locals()))

        #生成可信度标记脚本
        with open(confidence_shell,'w') as odata:
            odata.write(template)
        print('Mark confidence scripts have been generated to {}'.format(
            confidence_shell
        ))
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""
        add tags is used to handler add mutated samples or
        mark sites' confidence.
        --mapdir/-m is the path to bams and your bams should
        be stored in the subfile named by your corresponding
        samples' name,for example sample A's bam path is
        /path/to/mapping/A/A.final.bam
        """
    )
    parser.add_argument('-i','--infile',help='annovar infile')
    parser.add_argument('-o','--opre',help='output file prefix')
    
    parser.add_argument('-s','--mutated_samples',action="store_true",help='add mutated samples')
    parser.add_argument('-m','--mapdir',default=None,help='specify mapping directory if you want to output bam2html command or make confidence mark')
    parser.add_argument('-b','--bam2html',action="store_true",help='output bam2html command,--mapdir/-m is required')
    parser.add_argument('-c','--confidence',help="make confidence check,and you need to supply sample info,--mapdir/-m is required")
    parser.add_argument('-a','--acmg',help="add acmg tag to sites,value should be the path to ACMG directory")
    parser.add_argument('-n','--no_counter',action="store_true",help="do not add mutated samples count")
    
    args = vars(parser.parse_args())
    annovar_parser = AnnovarParser(args)
    annovar_parser.main()